{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-12-11T08:44:21.664117Z",
     "start_time": "2024-12-11T08:44:18.451894Z"
    }
   },
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import matplotlib.pyplot as plt\n",
    "from torch import nn\n",
    "from torch import Tensor\n",
    "from PIL import Image\n",
    "import torchvision\n",
    "from torchvision import transforms, models, datasets\n",
    "from torchvision.transforms import Compose, Resize, ToTensor\n",
    "import time\n",
    "from sklearn.model_selection import train_test_split\n",
    "import pickle"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-11T08:44:52.135198Z",
     "start_time": "2024-12-11T08:44:25.636430Z"
    }
   },
   "cell_type": "code",
   "source": [
    "print(os.listdir(\"./input\"))\n",
    "\n",
    "# df=pd.read_pickle(\"./input/WM811K.pkl\")\n",
    "df = pd.read_pickle(\"./input/LSWMD.pkl\")\n",
    "# info()函数用于打印DataFrame的简要摘要，显示有关DataFrame的信息，包括索引的数据类型dtype和列的数据类型dtype，非空值的数量和内存使用情况\n",
    "df.info()\n",
    "print(df.info())\n",
    "# head( )是指取数据的前n行数据，默认是前5行\n",
    "df.head()\n",
    "print(df.head())"
   ],
   "id": "50eee0ab28c0f433",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['save_weights', 'val', 'nonlabel', 'LSWMD.pkl', 'pretrained_weights', 'train', 'onpattern']\n",
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 811457 entries, 0 to 811456\n",
      "Data columns (total 6 columns):\n",
      " #   Column          Non-Null Count   Dtype  \n",
      "---  ------          --------------   -----  \n",
      " 0   waferMap        811457 non-null  object \n",
      " 1   dieSize         811457 non-null  float64\n",
      " 2   lotName         811457 non-null  object \n",
      " 3   waferIndex      811457 non-null  float64\n",
      " 4   trianTestLabel  811457 non-null  object \n",
      " 5   failureType     811457 non-null  object \n",
      "dtypes: float64(2), object(4)\n",
      "memory usage: 37.1+ MB\n",
      "<class 'pandas.core.frame.DataFrame'>\n",
      "RangeIndex: 811457 entries, 0 to 811456\n",
      "Data columns (total 6 columns):\n",
      " #   Column          Non-Null Count   Dtype  \n",
      "---  ------          --------------   -----  \n",
      " 0   waferMap        811457 non-null  object \n",
      " 1   dieSize         811457 non-null  float64\n",
      " 2   lotName         811457 non-null  object \n",
      " 3   waferIndex      811457 non-null  float64\n",
      " 4   trianTestLabel  811457 non-null  object \n",
      " 5   failureType     811457 non-null  object \n",
      "dtypes: float64(2), object(4)\n",
      "memory usage: 37.1+ MB\n",
      "None\n",
      "                                            waferMap  dieSize lotName  \\\n",
      "0  [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,...   1683.0    lot1   \n",
      "1  [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,...   1683.0    lot1   \n",
      "2  [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,...   1683.0    lot1   \n",
      "3  [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,...   1683.0    lot1   \n",
      "4  [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,...   1683.0    lot1   \n",
      "\n",
      "   waferIndex trianTestLabel failureType  \n",
      "0         1.0   [[Training]]    [[none]]  \n",
      "1         2.0   [[Training]]    [[none]]  \n",
      "2         3.0   [[Training]]    [[none]]  \n",
      "3         4.0   [[Training]]    [[none]]  \n",
      "4         5.0   [[Training]]    [[none]]  \n"
     ]
    }
   ],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-11T08:47:09.834142Z",
     "start_time": "2024-12-11T08:46:35.723098Z"
    }
   },
   "cell_type": "code",
   "source": [
    "df.tail()\n",
    "df['failureNum'] = df.failureType\n",
    "df['trainTestNum'] = df.trianTestLabel\n",
    "mapping_type = {'Center': 0, 'Donut': 1, 'Edge-Loc': 2, 'Edge-Ring': 3, 'Loc': 4, 'Random': 5, 'Scratch': 6,\n",
    "                'Near-full': 7, 'none': 8}\n",
    "mapping_traintest = {'Training': 0, 'Test': 1}\n",
    "df = df.replace({'failureNum': mapping_type, 'trainTestNum': mapping_traintest})\n",
    "# typelist = ['Center','Donut','Edge-Loc','Edge-Ring','Loc','Random','Scratch','Near-full','none']\n",
    "# df_withlabel = df[(df[\"failureNum\"] <= 8) & (df[\"failureNum\"] >= 0)]\n",
    "# df_withlabel =df_withlabel.reset_index()\n",
    "df_withpattern = df[(df[\"failureNum\"] <= 7) & (df[\"failureNum\"] >= 0)]\n",
    "df_withpattern = df_withpattern.reset_index()\n",
    "df_nonpattern = df[(df[\"failureNum\"] == 8)]"
   ],
   "id": "f65c09deeb2f7a33",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-11T08:47:55.627306Z",
     "start_time": "2024-12-11T08:47:55.621630Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def ensure_dirs_exist():\n",
    "    dirs = ['Center', 'Donut', 'Edge-Loc', 'Edge-Ring', 'Loc', 'Random', 'Scratch', 'Near-full','none']\n",
    "    train_dir = './input/train/'\n",
    "    val_dir = './input/val/'\n",
    "    for dir in dirs:\n",
    "        temp_dir = train_dir + dir\n",
    "        if not os.path.exists(temp_dir):\n",
    "            os.makedirs(temp_dir)\n",
    "        temp_dir = val_dir + dir\n",
    "        if not os.path.exists(temp_dir):\n",
    "            os.makedirs(temp_dir)"
   ],
   "id": "85a04917d972b998",
   "outputs": [],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-11T08:47:56.563053Z",
     "start_time": "2024-12-11T08:47:56.558483Z"
    }
   },
   "cell_type": "code",
   "source": "ensure_dirs_exist()",
   "id": "2e37413615c4aa5d",
   "outputs": [],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-11T08:48:16.350656Z",
     "start_time": "2024-12-11T08:47:59.966029Z"
    }
   },
   "cell_type": "code",
   "source": [
    "filepath = \"./input/train/\"\n",
    "for i in range(0, df_withpattern.shape[0]):\n",
    "    index = df_withpattern.index[i]\n",
    "    img = df_withpattern.waferMap[index]\n",
    "    label = str(df_withpattern.failureType[index][0]).replace(\"['\", \"\").replace(\"']\", \"\")\n",
    "    traintestlabel = df_withpattern.trianTestLabel[index]\n",
    "    if traintestlabel == 'Training':\n",
    "        plt.imsave(\"./input/train/\" + str(label) + \"/\" + str(label) + \"_\" + str(index) + \".jpg\", img)\n",
    "    else:\n",
    "        plt.imsave(\"./input/val/\" + str(label) + \"/\" + str(label) + \"_\" + str(index) + \".jpg\", img)\n",
    "    # img.save(\"./input/data/\" + str(label) + \"/\" + str(i) + \".jpg\")\n"
   ],
   "id": "f3584dd0680b7acf",
   "outputs": [],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-11T08:50:45.012133Z",
     "start_time": "2024-12-11T08:49:10.844088Z"
    }
   },
   "cell_type": "code",
   "source": [
    "filepath = \"./input/onpattern/\"\n",
    "for i in range(0, df_nonpattern.shape[0]):\n",
    "    index = df_nonpattern.index[i]\n",
    "    img = df_nonpattern.waferMap[index]\n",
    "    label = str(df_nonpattern.failureType[index][0]).replace(\"['\", \"\").replace(\"']\",\"\")\n",
    "    traintestlabel = df_nonpattern.trianTestLabel[index]\n",
    "    plt.imsave(filepath + str(index) + \".jpg\", img)"
   ],
   "id": "9b275d88c662aa6d",
   "outputs": [],
   "execution_count": 7
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-11T08:58:21.495105Z",
     "start_time": "2024-12-11T08:51:13.836453Z"
    }
   },
   "cell_type": "code",
   "source": [
    "filepath = \"./input/nonlabel/\"\n",
    "for i in range(0, df.shape[0]):\n",
    "    index = df.index[i]\n",
    "    img = df.waferMap[index]\n",
    "    type = df.failureType[index]\n",
    "    if not df.failureType[index]:\n",
    "        # label = str(df_nonlabel.failureType[index][0]).replace(\"['\", \"\").replace(\"']\",\"\")\n",
    "        plt.imsave(filepath + str(index) + \".jpg\", img)"
   ],
   "id": "477d8c0ee0c091e2",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_382354/2289993826.py:6: DeprecationWarning: The truth value of an empty array is ambiguous. Returning False, but in future this will result in an error. Use `array.size > 0` to check that an array is not empty.\n",
      "  if not df.failureType[index]:\n"
     ]
    }
   ],
   "execution_count": 8
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "f22df39a527fa8b2"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
